#include "main/utils.h"
using namespace std;

class Solution {
  public:
    int quick_select(vector<int>& nums, int left, int right, int k) {
      if (left == right) {
        return nums[k];
      }
      int partition = nums[left];
      int l = left - 1, r = right + 1;
      while (l < r) {
        do ++l; while (nums[l] < partition);
        do --r; while (nums[r] > partition);
        if (l < r) {
          swap(nums[l], nums[r]);
        }
      }
      if (r < k) {
        return quick_select(nums, r + 1, right, k);
      } else {
        return quick_select(nums, left, r, k);
      }
    }
    int findKthLargest(vector<int>& nums, int k) {
      int len = nums.size();
      return quick_select(nums, 0, len - 1, len - k);
    }

    void adjustHeap(vector<int>& nums, int start, int end) {
      int left = 2 * start + 1, right = 2 * start + 2, largest = start;
      if (left < end && nums[left] > nums[largest]) {
        largest = left;
      }
      if (right < end && nums[right] > nums[largest]) {
        largest = right;
      }
      if (largest != start) {
        swap(nums[start], nums[largest]);
        adjustHeap(nums, largest, end);
      }
    }
    void buildHeap(vector<int>& nums, int len) {
      for (int i = len / 2 - 1; i >= 0; --i) {
        adjustHeap(nums, i, len);
      }
    }
    int findKthLargest1(vector<int>& nums, int k) {
      int len = nums.size();
      buildHeap(nums, len);
      for (int i = len - 1; i >= nums.size() - k + 1; --i) {
        swap(nums[0], nums[i]);
        --len;
        adjustHeap(nums, 0, len);
      }
      return nums[0];
    }
};

int main() {
  vector<int> input = {3, 2, 1, 5, 6, 4};
  int k = 2;
  Solution su;
  cout << su.findKthLargest(input, k) << endl;
  cout << su.findKthLargest1(input, k) << endl;
}

